#! /usr/bin/env python
# coding=utf-8




pep_raw = r'pep_raw_ali.fa'
mRNA_raw = r'mRNA_raw_ali.fa'
rate_file = r'rate4site_out2'

new_pep_file = r'pep_cut_ali.fa'
new_mRNA_file = r'mRNA_cut_ali.fa'
cut_loci_file = r'cut_loci'

# 这里获得的列表的意思 为 1 1   33  M   0.1133  36/37
# 第一列 cut完的列
# 第二列 原先第一条序列的列号
# 第三列 原先比对矩阵的列号
# 第四列 原先第一条序列的氨基酸
# 第五列 rate4site 分值
# 第六列 占比物种



import os


from Bio import SeqIO
import re
import argparse
import sys



parser = argparse.ArgumentParser(
    description='''读取rate4site结果 根据cutoff 和least_nub值 进行切割并过滤
    用法:
    conver01_read_rate4site_for_every_ogs.py  -i OGS_for_testing -c 2.5 -l 30
    由大天才于2021年11月11日创建于浙江农业大学''')





parser.add_argument('-i',
                help='必须给定，OGS的工作目录')

parser.add_argument('-c',
                help='cutoff值,必须给定')

parser.add_argument('-l',
                help='最小比对数目值，必须给定')

args = parser.parse_args()

if not args.i or not args.c or not args.l:
    parser.print_help()
    sys.exit()




infile = args.i

cutoff = float(args.c)

least_nub  = int(args.l)



site_nub = 0
raw_nub = 0

nub_size = 0

for i in os.listdir(infile):
    nub_size += 1
    target_path = infile+'/'+i+'/'
    seq_pep = {str(j.name):str(j.seq) for j in SeqIO.parse(target_path+pep_raw,'fasta')}
    first_line = str(next(SeqIO.parse(target_path+pep_raw,'fasta')).seq)
    seq_mRNA= {str(j.name):str(j.seq) for j in  SeqIO.parse(target_path+mRNA_raw,'fasta')}

    first_one_lista = []
    socre_dic = {}


    # 读取 rate4site
    with open(target_path+rate_file) as fila:
        for j in fila:
            j = j.strip()
            if j!='' and j[0] != '#':
                k = re.split('[ ]+',j)
                if len(k) >= 5:
                    if float(k[2]) <= cutoff and int(k[5].split('/')[0])>=least_nub:
                        first_one_lista.append(int(k[0]))
                        socre_dic[int(k[0])] = (k[1],k[2],k[5])

    # 将rate4site的结果转化为 真实比对的位置
    #seq_pep[]
    true_loci_dic = {}
    n = 0
    raw_nub+= len(first_line)
    for x,y in enumerate(first_line):
        if y != '-':
            n += 1
            true_loci_dic[n] = x
    new_loci_dic = {}
    new_pep_dic = {j:'' for j in seq_pep}
    new_mRNA_dic = {j:'' for j in seq_pep}
    n = 0
    for j in first_one_lista:
        n+=1
        new_loci = true_loci_dic[j]
        new_loci_dic[n] = (n,j, true_loci_dic[j]+1,socre_dic[j][0],socre_dic[j][1],socre_dic[j][2])
        for x in seq_pep:
            new_pep_dic[x] += seq_pep[x][new_loci]
        for x in seq_mRNA:
            new_mRNA_dic[x] += seq_mRNA[x][new_loci*3:new_loci*3+3]
    new_pep=open(target_path+new_pep_file,'w')
    new_mRNA=open(target_path+new_mRNA_file,'w')
    cut_loci = open(target_path + cut_loci_file,'w')
    #print(new_mRNA_dic)
    for j in new_mRNA_dic:
        new_mRNA.write('>'+j+'\n'+new_mRNA_dic[j]+'\n')
    for j in new_pep_dic:
        new_pep.write('>'+j+'\n'+new_pep_dic[j]+'\n')

    site_nub += len(new_pep_dic[j])

    cut_loci.write('\n'.join(['\t'.join([str(x) for x in new_loci_dic[n]]) for n in new_loci_dic]))
    cut_loci.close()
    new_pep.close()
    new_mRNA.close()
    #print(new_loci_dic)
print('共有 '+str(nub_size)+' 个OGS')
print('原先有 ' + str(raw_nub)+' 个AA sites 含有gap')
print('经过过滤 还剩下 ' + str(site_nub)+' 个AA sites')

    

    #print(seq_pep)

